# !/usr/bin/env python
import cv2 as cv
import numpy as np

SZ = 40
bin_n = 16  # Number of bins
affine_flags = cv.WARP_INVERSE_MAP | cv.INTER_LINEAR


def deskew(img):
	m = cv.moments(img)
	if abs(m['mu02']) < 1e-2:
		return img.copy()
	skew = m['mu11'] / m['mu02']
	M = np.float32([[1, skew, -0.5 * SZ * skew], [0, 1, 0]])
	img = cv.warpAffine(img, M, (SZ, SZ), flags=affine_flags)
	return img


def hog(img):
	gx = cv.Sobel(img, cv.CV_32F, 1, 0)
	gy = cv.Sobel(img, cv.CV_32F, 0, 1)
	mag, ang = cv.cartToPolar(gx, gy)
	bins = np.int32(bin_n * ang / (2 * np.pi))  # quantizing binvalues in (0...16)
	# print(bins.shape)
	bin_cells = bins[:20, :20], bins[20:, :20], bins[:20, 20:], bins[20:, 20:]
	mag_cells = mag[:20, :20], mag[20:, :20], mag[:20, 20:], mag[20:, 20:]
	hists = [np.bincount(b.ravel(), m.ravel(), bin_n) for b, m in zip(bin_cells, mag_cells)]
	hist = np.hstack(hists)  # hist is a 64 bit vector
	return hist


img = cv.imread('digits.jpg', 0)
if img is None:
	raise Exception("we need the digits.png image from samples/data here !")
cells = [np.hsplit(row, 40) for row in np.vsplit(img, 40)]
# First half is trainData, remaining is testData
train_cells = [i[:20] for i in cells]
test_cells = [i[20:] for i in cells]
deskewed = [list(map(deskew, row)) for row in train_cells]
hogdata = [list(map(hog, row)) for row in deskewed]
trainData = np.float32(hogdata).reshape(-1, 64)
responses = np.repeat(np.arange(8) + 1, 100)[:, np.newaxis]
svm = cv.ml.SVM_create()
svm.setKernel(cv.ml.SVM_LINEAR)
svm.setType(cv.ml.SVM_C_SVC)
svm.setC(2.67)
svm.setGamma(5.383)
svm.train(trainData, cv.ml.ROW_SAMPLE, responses)
svm.save('svm_data.dat')
deskewed = [list(map(deskew, row)) for row in test_cells]
hogdata = [list(map(hog, row)) for row in deskewed]
testData = np.float32(hogdata).reshape(-1, bin_n * 4)
result = svm.predict(testData)[1]
print(result)

mask = result == responses
correct = np.count_nonzero(mask)
print(correct * 100.0 / result.size)

